{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这四个台风为一组，左边放Umax RMSE，右边放distance RMSE\n",
    "然后后面四个台风为一组，也是左边放Umax RMSE，右边放distance RMSE\n",
    "左列公用一个colorbar,右列公用一个colorbar,色标放在图片的最下面，并排放\n",
    "Y轴时间太密了，改为间隔两天标注，不用写年份，上方不用写台风名字，台风名字统一放在每张图的左下角空白处，写的大一些，例如DUKUSRI（2023）\n",
    "\n",
    "\n",
    "Figure3.1:  dusurui gaemi haikui kangni \n",
    "Figure3.2:  koinu mojie saola shantuo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os,glob,math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.lines as mlines\n",
    "import matplotlib.dates as mdates\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter\n",
    "import cartopy.feature as cfeature\n",
    "import shapely.geometry as sgeom\n",
    "from datetime import datetime,timedelta\n",
    "from tqdm import tqdm\n",
    "from global_land_mask import globe\n",
    "from typlot.scripts.gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall\n",
    "from geopy.distance import geodesic\n",
    "import matplotlib.ticker as ticker\n",
    "import seaborn as sns\n",
    "import xarray as xr\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from typlot.config.global_config import *\n",
    "from geopy.distance import geodesic\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "\n",
    "ini_time_mode = ['00','12']\n",
    "# names = names[:4]\n",
    "names = names[4:]\n",
    "years = years[4:]\n",
    "\n",
    "# names = ['dusurui_16']\n",
    "tynames,tyids = split_str_id(names)\n",
    "obs_baseline='land'  # ‘land’ 'RI'\n",
    "RIstd = 7\n",
    "tyrmse = {}\n",
    "track_id = np.arange(1,52)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  制作数据集\n",
    "\n",
    "dalist = []\n",
    "tyobslist = []\n",
    "\n",
    "for ty,tyid in zip(tynames,tyids):  \n",
    "    # 初始化\n",
    "    directory  = os.path.join(global_ensdir,f'{ty}_{tyid}')\n",
    "    dates_name = sorted(os.listdir(directory))\n",
    "    dates_name = [i for i in dates_name if i[-2:] in ini_time_mode ]\n",
    "    obs_path   = os.path.join( global_obsdir, f'{ty}_CMAobs.txt')\n",
    "    tyobs = tydat_CMA(obs_path)\n",
    "    variable = ['umax','pmin','dist']\n",
    "    print(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc'))\n",
    "    if not os.path.exists(ncpath:=os.path.join(basedir,f'{ty}_pmin_umax_dist.nc')):\n",
    "        print(ncpath,'not found')\n",
    "        #### 计算最长的leadtime，确定后用来制作DataArray\n",
    "        leadtime_list=[]\n",
    "        # 遍历起报时间\n",
    "        for date_name in tqdm(dates_name,total=len(dates_name),desc='Filtering max leadtime'):\n",
    "            dir_date = os.path.join(directory,date_name)\n",
    "            name_ensembles = os.listdir(dir_date)\n",
    "            path_ensembles = [os.path.join(dir_date,i) for i in name_ensembles if i.startswith('TRACK')]\n",
    "            # 遍历集合\n",
    "            for path in path_ensembles:\n",
    "                t_id = int(path.split('_')[-1])\n",
    "                if t_id==0:\n",
    "                    continue # ERA5 情况\n",
    "                else:\n",
    "                    tyens = tydat(path,RIstd)\n",
    "                    ini_time = datetime.strptime(date_name, '%Y%m%d%H')\n",
    "                    end_time = tyens.time[-1]\n",
    "                    leadtime_list.append(end_time-ini_time)\n",
    "        \n",
    "        max_leadtime = max(leadtime_list)\n",
    "        n_leadtime = np.ceil(max_leadtime / timedelta(hours=6)).astype(int)+1\n",
    "        \n",
    "        #### 创建DataArray  dims=['variable','start_time','lead_time']\n",
    "        lead_time = 6*np.arange(n_leadtime)  # 6 for 6h/index.\n",
    "        da = xr.DataArray(np.nan,dims=['variable','start_time','lead_time'],\n",
    "                    coords={'variable':variable ,'start_time':dates_name ,'lead_time':lead_time})\n",
    "\n",
    "        #### 开始填充da,逐个元素计算\n",
    "        # 遍历start_time\n",
    "        for date_name in tqdm(dates_name,total=len(dates_name),desc='Calcing RMSE on members'):\n",
    "            dir_date = os.path.join(directory,date_name)\n",
    "            name_ensembles = os.listdir(dir_date)\n",
    "            path_ensembles = [os.path.join(dir_date,i) for i in name_ensembles if i.startswith('TRACK')]\n",
    "            # 遍历 lead_time\n",
    "            for i in lead_time:\n",
    "                v_ens_list,p_ens_list = [],[]\n",
    "                v_obs_list,p_obs_list = [],[]\n",
    "                lat_ens_list,lat_obs_list = [],[]\n",
    "                lon_ens_list,lon_obs_list = [],[]\n",
    "                leadtime = datetime.strptime(date_name,'%Y%m%d%H') + timedelta(hours=int(i))\n",
    "                # 遍历所有的成员，计算\n",
    "                for path in path_ensembles:\n",
    "                    t_id = int(path.split('_')[-1])\n",
    "                    if t_id==0:\n",
    "                        continue # 再分析 情况\n",
    "                    else:\n",
    "                        tyens = tydat(path,RIstd)\n",
    "                        # 只留下leadtime同时存在于ens和obs的情况\n",
    "                        if (leadtime in tyobs.time) and (leadtime in tyens.time):\n",
    "                            v_ens_list.append(tyens.umax[tyens.time==leadtime][0])\n",
    "                            v_obs_list.append(tyobs.umax[tyobs.time==leadtime][0])\n",
    "                            p_ens_list.append(tyens.pmin[tyens.time==leadtime][0])\n",
    "                            p_obs_list.append(tyobs.pmin[tyobs.time==leadtime][0])\n",
    "                            lat_ens_list.append(tyens.lat[tyens.time==leadtime][0])\n",
    "                            lat_obs_list.append(tyobs.lat[tyobs.time==leadtime][0])\n",
    "                            lon_ens_list.append(tyens.lon[tyens.time==leadtime][0])\n",
    "                            lon_obs_list.append(tyobs.lon[tyobs.time==leadtime][0])\n",
    "                        else:\n",
    "                            v_ens_list.append(np.nan)\n",
    "                            v_obs_list.append(np.nan)\n",
    "                            p_ens_list.append(np.nan)\n",
    "                            p_obs_list.append(np.nan)\n",
    "                            lat_ens_list.append(np.nan)\n",
    "                            lat_obs_list.append(np.nan)\n",
    "                            lon_ens_list.append(np.nan)\n",
    "                            lon_obs_list.append(np.nan)\n",
    "                v_ens,p_ens,v_obs,p_obs = np.array(v_ens_list),np.array(p_ens_list),np.array(v_obs_list),np.array(p_obs_list)\n",
    "                lat_ens,lat_obs,lon_ens,lon_obs = np.array(lat_ens_list),np.array(lat_obs_list),np.array(lon_ens_list),np.array(lon_obs_list)\n",
    "                dist_rmse_list = [geodesic((i,j),(k,l)).kilometers if not math.isnan(i*j*k*l) else np.nan for i,j,k,l in zip(lat_ens,lon_ens,lat_obs,lon_obs) ]\n",
    "                \n",
    "                #### 计算该leadtime的RMSE\n",
    "                umax_rmse = np.sqrt(np.nanmean((v_ens-v_obs)**2))\n",
    "                pmin_rmse = np.sqrt(np.nanmean((p_ens-p_obs)**2))\n",
    "                dist_rmse = np.sqrt(np.nanmean((np.array(dist_rmse_list))**2))\n",
    "\n",
    "                #### 补充DataArray\n",
    "                da.loc['umax',date_name,i] = umax_rmse\n",
    "                da.loc['pmin',date_name,i] = pmin_rmse\n",
    "                da.loc['dist',date_name,i] = dist_rmse\n",
    "\n",
    "        da.to_netcdf(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc'))\n",
    "    else :\n",
    "        da = xr.open_dataarray(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc'))\n",
    "    \n",
    "    tyobslist.append(tyobs)\n",
    "    dalist.append(da)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all for red dot indexes\n",
    "\n",
    "def get_RI_init_time(tyobs, da):\n",
    "    \"\"\"\n",
    "    根据tyobs的信息，找到观测中的RI时刻。\n",
    "    但是观测RI往往不在严格的00，12h这种起报时间上。\n",
    "    因此要把观测RI转换到最近的RI时刻上去。\n",
    "    \"\"\"\n",
    "    from datetime import datetime\n",
    "    index = count_rapidgrow(7, tyobs.umax, tyobs.time).astype(bool)\n",
    "    obs_RI_moments = tyobs.time[index]\n",
    "    # 将 da.start_time 从字符串转换为 datetime 对象\n",
    "    da_start_times = [datetime.strptime(str(t), \"%Y%m%d%H\") for t in da.start_time.values]\n",
    "    # 找到da中最近的RI起报时刻\n",
    "    da_RI_moments = []\n",
    "    for obs_time in obs_RI_moments:\n",
    "        # 计算时间差\n",
    "        time_diffs = np.array([(dt - obs_time).total_seconds() for dt in da_start_times])\n",
    "        time_diffs[time_diffs < 0] = np.inf\n",
    "        # 找到最小正数时间差的索引\n",
    "        min_idx = np.argmin(time_diffs)\n",
    "        # 添加对应的 start_time 字符串\n",
    "        da_RI_moments.append(da.start_time.values[min_idx])\n",
    "    return da_RI_moments\n",
    "\n",
    "\n",
    "def ceil_timedelta(td):\n",
    "    \"\"\"将 Timedelta 向上取整到天\"\"\"\n",
    "    return pd.Timedelta(days=int(np.ceil(td.total_seconds() / 86400)))\n",
    "\n",
    "\n",
    "\n",
    "def gen_RI_scatter(tyobs,da):\n",
    "    # 根据da来得到y和x轴的时间维度信息\n",
    "    init_times = [pd.Timestamp(i+'00') for i in da.start_time.values]\n",
    "    dy = init_times[1]-init_times[0]\n",
    "    dx = pd.Timedelta(da.lead_time.values[1]-da.lead_time.values[0],'h')\n",
    "\n",
    "    RI_ini_times = [pd.Timestamp(i+'00') for i in get_RI_init_time(tyobs, da)]\n",
    "    x = []\n",
    "    y = []\n",
    "\n",
    "    for RI_ini_time in RI_ini_times:\n",
    "        ini_time = RI_ini_time\n",
    "        # change me here to 15d lead time \n",
    "        while  RI_ini_time-ini_time <= pd.Timedelta(da.lead_time.values[-1],'h')  and  \\\n",
    "            ini_time>= init_times[0] :\n",
    "            x.append((RI_ini_time-ini_time))\n",
    "            y.append(ini_time)\n",
    "            ini_time -= dy\n",
    "    \n",
    "    # 对x进行处理, 向上取整到天\n",
    "    # x = [ceil_timedelta(xx) for xx in x]\n",
    "\n",
    "    # 转换为index\n",
    "    yindex = np.searchsorted(init_times,np.array(y))\n",
    "    lead_time_timedelta = [pd.Timedelta(i,'h') for i in da.lead_time.values]\n",
    "    xindex = np.searchsorted(lead_time_timedelta,np.array(x))\n",
    "    return yindex,xindex\n",
    "\n",
    "\n",
    "def split_yx(y, x):\n",
    "    ''' \n",
    "    直接生成的yx包含多个RI线，这里将其根据时间的不同拆分\n",
    "    Args:\n",
    "        y: 直接生成的y轴时间index\n",
    "        x: 直接生成的x轴时间index\n",
    "    Returns:\n",
    "        split_data: 一个列表，每个元素是一个包含(y, x)对的元组，表示每个时间段的数据。\n",
    "    '''\n",
    "    split_data = []  # 存储拆分后的数据\n",
    "    current_group = []  # 当前组的点\n",
    "\n",
    "    for i in range(len(x)):\n",
    "        # 检查 x 是否为 0，表示一个新的组开始\n",
    "        if x[i] == 0 and current_group:\n",
    "            # 如果当前组不为空，保存当前组并开始新的组\n",
    "            split_data.append(current_group)\n",
    "            current_group = []  # 重置当前组\n",
    "\n",
    "        # 将 (y, x) 对添加到当前组\n",
    "        current_group.append((y[i], x[i]))\n",
    "\n",
    "    # 添加最后一组（如果存在）\n",
    "    if current_group:\n",
    "        split_data.append(current_group)\n",
    "\n",
    "    return split_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# ========== 1️⃣ 创建带 colorbar 空间的布局 ==========\n",
    "fig = plt.figure(figsize=(15, 20))  # 稍微增加高度，为 colorbar 留空间\n",
    "\n",
    "# 使用 GridSpec 创建布局：\n",
    "# - 上面 len(names) 行用于子图\n",
    "# - 下面 1 行用于 colorbar\n",
    "gs = gridspec.GridSpec(\n",
    "    len(names) + 1,  # 总行数 = 子图行数 + colorbar 行\n",
    "    2,               # 2列\n",
    "    height_ratios=[1]*len(names) + [0.08],  # 前面的行高度相同，最后一行（colorbar）高度为 0.08\n",
    "    hspace=0.25,      # 行间距\n",
    "    wspace=0.25      # 列间距\n",
    ")\n",
    "\n",
    "# 创建子图数组\n",
    "axes = np.array([[fig.add_subplot(gs[i, j]) for j in range(2)] for i in range(len(names))])\n",
    "\n",
    "# ========== 2️⃣ 绘制热图（保持你的原始逻辑）==========\n",
    "im_umax = None  # 用于存储 umax 的 mappable\n",
    "im_dist = None  # 用于存储 dist 的 mappable\n",
    "\n",
    "for j, (da, tyobs,name, year) in enumerate(zip(dalist, tyobslist, tynames, years)):\n",
    "    for i, var in enumerate(['umax', 'dist']):\n",
    "        \n",
    "        ax = axes[j, i]\n",
    "\n",
    "        # 1) 读取数据               \n",
    "        data = da.loc[var].values.astype(float)\n",
    "        data = np.where(data > 10000, np.nan, data)\n",
    "        data = np.ma.masked_invalid(data)\n",
    "\n",
    "        # 2) 坐标 \n",
    "        members = da.lead_time.values  \n",
    "        times   = da.start_time.values \n",
    "\n",
    "        # 🔧 预报时效转换\n",
    "        lead_days = members / 24\n",
    "        day_indices = []\n",
    "        day_labels = []\n",
    "\n",
    "        for idx, day in enumerate(lead_days):\n",
    "            if day.is_integer():\n",
    "                day_indices.append(idx)\n",
    "                day_labels.append(f'{int(day)}d')\n",
    "\n",
    "        if not day_indices:\n",
    "            step = 24\n",
    "            day_indices = list(range(0, len(members), step))\n",
    "            day_labels = [f'{int(members[idx]/24)}d' for idx in day_indices]\n",
    "            \n",
    "        # ⭐ 根据变量类型设置颜色范围\n",
    "        if var == 'pmin':\n",
    "            vmin, vmax = 0, 50\n",
    "            cbar_label = 'RMSE (hPa)'\n",
    "        elif var == 'umax':\n",
    "            vmin, vmax = 0, 50\n",
    "            cbar_label = 'RMSE (m/s)'\n",
    "        elif var == 'dist':    \n",
    "            vmin, vmax = 0, 3000\n",
    "            cbar_label = 'RMSE (km)'\n",
    "        else:\n",
    "            vmin, vmax = None, None\n",
    "\n",
    "        # ✅ 绘制热图\n",
    "        im = ax.imshow(\n",
    "            data,                    \n",
    "            aspect='auto',\n",
    "            interpolation='none',\n",
    "            cmap='viridis',\n",
    "            origin='lower',\n",
    "            vmin=vmin,\n",
    "            vmax=vmax\n",
    "        )\n",
    "\n",
    "        # 🔧 保存最后一个 mappable（用于创建 colorbar）\n",
    "        if var == 'umax':\n",
    "            im_umax = im\n",
    "        elif var == 'dist':\n",
    "            im_dist = im\n",
    "\n",
    "        # 🔧 设置刻度\n",
    "        ax.set_yticks(np.arange(len(times))[::4])\n",
    "        yticklabels = [i[4:] for i in times[::4]]\n",
    "        ax.set_yticklabels(yticklabels, rotation=0, fontsize=8)\n",
    "        ax.set_xticks(day_indices)\n",
    "        ax.set_xticklabels(day_labels, fontsize=8)\n",
    "\n",
    "        # 🔧 设置标题\n",
    "        varnamedict={'umax':'Maximum wind speed RMSE','dist':'Position RMSE'}\n",
    "\n",
    "        if j == 0:\n",
    "            ax.set_title(f'{varnamedict[var]}', fontsize=15, fontweight='bold')\n",
    "\n",
    "        # 🔧 散点图\n",
    "        y, x = gen_RI_scatter(tyobs, da)\n",
    "        groups = split_yx(y, x)\n",
    "        # 绘制每个组的折线图\n",
    "        for group in groups:\n",
    "            y_group, x_group = zip(*group)\n",
    "            ax.plot(x_group, y_group, color='red', linewidth=2, linestyle='-', alpha=0.7)  # 添加标签以便图例\n",
    "\n",
    "        #  底部标签\n",
    "        ax.text(\n",
    "            0.75, 0.88,\n",
    "            f\"{name}({year})\",\n",
    "            transform=ax.transAxes,\n",
    "            fontsize=10,\n",
    "            fontweight='bold',\n",
    "            color='white',\n",
    "            verticalalignment='bottom',\n",
    "            horizontalalignment='left',\n",
    "            bbox=dict(\n",
    "                boxstyle='round,pad=0.5',\n",
    "                edgecolor='none',\n",
    "                alpha=0.7\n",
    "            )\n",
    "        )\n",
    "\n",
    "\n",
    "# ========== 3️⃣ 在底部添加共享 colorbar ==========\n",
    "# 为 umax 添加 colorbar（左列）\n",
    "cax_umax = fig.add_subplot(gs[-1, 0])  # 最后一行，第一列\n",
    "cbar_umax = fig.colorbar(im_umax, cax=cax_umax, orientation='horizontal')\n",
    "cbar_umax.set_label('RMSE (m/s)', fontsize=12)\n",
    "cbar_umax.ax.tick_params(labelsize=10)\n",
    "\n",
    "# 为 dist 添加 colorbar（右列）\n",
    "cax_dist = fig.add_subplot(gs[-1, 1])  # 最后一行，第二列\n",
    "cbar_dist = fig.colorbar(im_dist, cax=cax_dist, orientation='horizontal')\n",
    "cbar_dist.set_label('RMSE (km)', fontsize=12 )\n",
    "cbar_dist.ax.tick_params(labelsize=10)\n",
    "\n",
    "# ========== 4️⃣ 保存/显示 ==========\n",
    "plt.tight_layout()\n",
    "\n",
    "# 在循环结束后添加标签\n",
    "labels = [\"(a)\",\"(b)\",\"(c)\",\"(d)\",\"(e)\",\"(f)\",\"(g)\",\"(h)\"]\n",
    "for ax, label in zip(axes.flatten(), labels):\n",
    "    ax.text(0.05, 0.1, label,\n",
    "            transform=ax.transAxes,  # 使用轴坐标系 (0-1)\n",
    "            fontsize=14,\n",
    "            fontweight='bold',\n",
    "            va='top',                # 垂直对齐：顶部\n",
    "            ha='left',              # 水平对齐：右侧\n",
    "            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
    "\n",
    "if draw_opt:\n",
    "    plt.show()\n",
    "else:\n",
    "    plt.savefig(\n",
    "        os.path.join(paperdir, f\"fig32.{pictype}\"),\n",
    "        dpi=600,\n",
    "        bbox_inches='tight'\n",
    "    )\n",
    "    plt.close() \n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wmq",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
